import random
from QKDNetwork import QKDNetwork
from compare import Compare
from docplex.mp.model import Model
from add_ps import AddPS
import networkx as nx

# 节点个数： 20-210          默认：60      迭代：20 + cnt * 10
# 源目对数： 10-200          默认：50      迭代：10 + cnt * 10
# 阿尔法值： 0.35-1.3        默认：0.2     迭代：0.2 + cnt * 0.05

def get_sps_add_ps_num(net: QKDNetwork, init_ps: list):
    sps_count = 0
    for node in net.G.nodes:
        if(node in init_ps):
            net.G.nodes[node]["transmitter"] = 1
    for sd in net.sd_list:
        path = nx.shortest_path(net.G, sd[0], sd[1])
        for path_index, path_node in enumerate(path):
            if(path_index % 2 == 0 and net.G.nodes[path_node]["transmitter"] == 0):
                net.G.nodes[path_node]["transmitter"] = 1
                sps_count += 1
    return sps_count

def get_ras_add_ps_num(net: QKDNetwork, init_ps: list):
    ras_count = 0
    for node in net.G.nodes:
        if(node in init_ps):
            net.G.nodes[node]["transmitter"] = 1
    for edge in net.G.edges:
        if(net.G.nodes[edge[0]]["transmitter"] == 0 and net.G.nodes[edge[1]]["transmitter"] == 0):
            ras_count += 1
            if(net.G.degree(edge[0]) > net.G.degree(edge[1])):
                net.G.nodes[edge[0]]["transmitter"] = 1
            else:
                net.G.nodes[edge[1]]["transmitter"] = 1
            
    for node in net.G.nodes:
        if(net.G.nodes[node]["transmitter"] == 1):
            ras_count += 1
    return ras_count

def clear_all_placed_ps(net: QKDNetwork):
    for node in net.G.nodes:
        if(net.G.nodes[node]["transmitter"] == 1):
            net.G.nodes[node]["transmitter"] = 0

def get_add_ps_num():
    q = QKDNetwork(showTopology=False, num_node=40, sd_num=40, alpha=0.2, no_add_ps=True)
    init_ps = random.sample(list(range(q.num_nodes)), 5)
    print("准备完成")
    print("RAS:", get_ras_add_ps_num(q, init_ps))
    clear_all_placed_ps(q)
    print("SPS:", get_sps_add_ps_num(q, init_ps))
    clear_all_placed_ps(q)
    ret = AddPS(q).run_expr(init_ps)
    print("APSP:", ret[0])
    print("ILP:", ret[1])

get_add_ps_num()
